import numpy as np
import cv2
import matplotlib.pyplot as plt
from scipy.fft import dct, idct
import os
from matplotlib import font_manager

# 设置中文字体
plt.rcParams['font.sans-serif'] = ['Arial Unicode MS']  # Mac系统的中文字体
plt.rcParams['axes.unicode_minus'] = False  # 解决负号显示问题

class MedicalImageProcessor:
    def __init__(self):
        self.input_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'images', 'input')
        self.output_dir = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'images', 'output')
        
        # 确保输出目录存在
        os.makedirs(self.output_dir, exist_ok=True)
        
        # 设置matplotlib的字体
        self._setup_font()

    def _setup_font(self):
        """设置matplotlib的字体"""
        # 尝试多个常见的中文字体
        chinese_fonts = ['Arial Unicode MS', 'Microsoft YaHei', 'SimHei', 'PingFang SC']
        font_set = False
        
        for font_name in chinese_fonts:
            try:
                font_path = font_manager.findfont(font_manager.FontProperties(family=font_name))
                if font_path:
                    plt.rcParams['font.sans-serif'] = [font_name]
                    font_set = True
                    break
            except:
                continue
                
        if not font_set:
            print("警告：未能找到合适的中文字体，文字可能显示为乱码")

    def pseudo_color(self, image_path):
        """医学图像伪彩处理"""
        # 读取图像
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")

        # 获取图像的高度和宽度
        h, w = img.shape
        L = 255
        img = img.astype(np.float32)  # 转换为浮点数进行计算

        # 预定义伪彩色通道
        r = np.zeros_like(img)
        g = np.zeros_like(img)
        b = np.zeros_like(img)

        # 使用条件索引映射
        r = np.where(
            img <= L / 4,
            0,
            np.where(
                img <= L / 2,
                0,
                np.where(img <= 3 * L / 4, np.clip(4 * img - 2 * L, 0, L), L),
            ),
        )
        g = np.where(
            img <= L / 4,
            np.clip(4 * img, 0, L),
            np.where(
                img <= L / 2,
                L,
                np.where(img <= 3 * L / 4, L, np.clip(-4 * img + 4 * L, 0, L)),
            ),
        )
        b = np.where(
            img <= L / 4, L, np.where(img <= L / 2, np.clip(-4 * img + 2 * L, 0, L), 0)
        )

        # 合并 R, G, B 通道形成伪彩色图像
        pseudo_color = np.stack([r, g, b], axis=-1).astype(np.uint8)
        
        # 保存结果
        output_path = os.path.join(self.output_dir, 'pseudo_color_result.jpg')
        cv2.imwrite(output_path, cv2.cvtColor(pseudo_color, cv2.COLOR_RGB2BGR))
        
        # 显示结果
        plt.figure(figsize=(12, 4))
        plt.subplot(121)
        plt.imshow(img.astype(np.uint8), cmap='gray')
        plt.title('原始图像')
        plt.axis('off')
        plt.subplot(122)
        plt.imshow(pseudo_color)
        plt.title('伪彩色处理结果')
        plt.axis('off')
        plt.savefig(os.path.join(self.output_dir, 'pseudo_color_comparison.jpg'))
        plt.close()

    def arnold_scramble(self, img, iterations=1):
        """Arnold置乱算法"""
        h, w = img.shape
        scrambled = img.copy()
        for _ in range(iterations):
            for i in range(h):
                for j in range(w):
                    x = (i + j) % h
                    y = (i + 2*j) % w
                    scrambled[i, j] = img[x, y]
        return scrambled

    def dct_watermark(self, image_path, watermark_text="MEDICAL"):
        """DCT水印嵌入"""
        # 读取图像
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")

        # 调整图像大小
        M = 1600
        img = cv2.resize(img, (M, M))

        # 创建水印图像
        N = 200
        watermark = np.zeros((N, N), dtype=np.uint8)
        cv2.putText(watermark, watermark_text, (50, 100), cv2.FONT_HERSHEY_SIMPLEX, 1, 255, 2)
        _, watermark = cv2.threshold(watermark, 0.65 * 255, 255, cv2.THRESH_BINARY)

        # 保存原始图像副本
        A = img.copy()

        # 置乱水印图像
        n = 10
        a = 3
        b = 5
        for i in range(n):
            imgn = np.zeros((N, N), dtype=np.uint8)
            for y in range(N):
                for x in range(N):
                    xx = (x + b * y) % N
                    yy = (a * x + (a * b + 1) * y) % N
                    imgn[yy, xx] = watermark[y, x]
            watermark = imgn

        # DCT水印嵌入
        K = 8
        alpha = 0.03
        for p in range(N):
            for q in range(N):
                x = p * K
                y = q * K
                BLOCK = img[x:x+K, y:y+K].astype(np.float32)
                BLOCK = cv2.dct(BLOCK)
                a = -1 if watermark[p, q] == 0 else 1
                BLOCK = BLOCK * (1 + a * alpha)
                BLOCK = cv2.idct(BLOCK)
                img[x:x+K, y:y+K] = BLOCK

        # 保存结果
        output_path = os.path.join(self.output_dir, 'watermarked_image.jpg')
        cv2.imwrite(output_path, img)

        # 显示结果
        plt.figure(figsize=(12, 4))
        plt.subplot(121)
        plt.imshow(A, cmap='gray')
        plt.title('原始图像')
        plt.axis('off')
        plt.subplot(122)
        plt.imshow(img, cmap='gray')
        plt.title('添加水印后的图像')
        plt.axis('off')
        plt.savefig(os.path.join(self.output_dir, 'watermark_comparison.jpg'))
        plt.close()

    def nearest_neighbor_interpolation(self, img, scale_factor):
        """最近邻插值"""
        h, w = img.shape
        new_h, new_w = int(h * scale_factor), int(w * scale_factor)
        resized = np.zeros((new_h, new_w), dtype=img.dtype)
        
        for i in range(new_h):
            for j in range(new_w):
                x = int(i / scale_factor)
                y = int(j / scale_factor)
                resized[i, j] = img[x, y]
        
        return resized

    def bilinear_interpolation(self, img, scale_factor):
        """双线性插值"""
        h, w = img.shape
        new_h, new_w = int(h * scale_factor), int(w * scale_factor)
        resized = np.zeros((new_h, new_w), dtype=img.dtype)
        
        for i in range(new_h):
            for j in range(new_w):
                x = i / scale_factor
                y = j / scale_factor
                
                x1, y1 = int(x), int(y)
                x2, y2 = min(x1 + 1, h - 1), min(y1 + 1, w - 1)
                
                # 计算权重
                wx = x - x1
                wy = y - y1
                
                # 双线性插值
                resized[i, j] = (1 - wx) * (1 - wy) * img[x1, y1] + \
                              wx * (1 - wy) * img[x2, y1] + \
                              (1 - wx) * wy * img[x1, y2] + \
                              wx * wy * img[x2, y2]
        
        return resized

    def image_interpolation(self, image_path, scale_factor=2.0):
        """图像插值处理"""
        # 读取图像
        img = cv2.imread(image_path, cv2.IMREAD_GRAYSCALE)
        if img is None:
            raise ValueError(f"无法读取图像: {image_path}")

        # 进行最近邻插值
        nn_result = self.nearest_neighbor_interpolation(img, scale_factor)
        
        # 进行双线性插值
        bilinear_result = self.bilinear_interpolation(img, scale_factor)

        # 保存结果
        cv2.imwrite(os.path.join(self.output_dir, 'nearest_neighbor_result.jpg'), nn_result)
        cv2.imwrite(os.path.join(self.output_dir, 'bilinear_result.jpg'), bilinear_result)

        # 显示结果
        plt.figure(figsize=(15, 5))
        plt.subplot(131)
        plt.imshow(img, cmap='gray')
        plt.title('原始图像')
        plt.axis('off')
        plt.subplot(132)
        plt.imshow(nn_result, cmap='gray')
        plt.title('最近邻插值结果')
        plt.axis('off')
        plt.subplot(133)
        plt.imshow(bilinear_result, cmap='gray')
        plt.title('双线性插值结果')
        plt.axis('off')
        plt.savefig(os.path.join(self.output_dir, 'interpolation_comparison.jpg'))
        plt.close()

def main():
    processor = MedicalImageProcessor()
    
    # 设置输入图像路径
    input_image = os.path.join(processor.input_dir, 'input.jpg')
    
    # 确保输入图像存在
    if not os.path.exists(input_image):
        print(f"请将输入图像放在 {processor.input_dir} 目录下，并命名为 input.jpg")
        return

    try:
        # 执行伪彩处理
        print("正在进行伪彩处理...")
        processor.pseudo_color(input_image)
        
        # 执行DCT水印
        print("正在进行DCT水印处理...")
        processor.dct_watermark(input_image)
        
        # 执行图像插值
        print("正在进行图像插值处理...")
        processor.image_interpolation(input_image)
        
        print("所有处理已完成！结果保存在 output 目录中。")
        
    except Exception as e:
        print(f"处理过程中出现错误: {str(e)}")

if __name__ == "__main__":
    main() 